{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c3dc1c7a",
   "metadata": {},
   "source": [
    "# 仿真量化测试"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79988179",
   "metadata": {},
   "source": [
    "定义前端网络"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7fa28eaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "class M(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv2d = torch.nn.Conv2d(3, 16, 3, 1, 1, bias=False)\n",
    "        self.relu = torch.nn.ReLU()\n",
    "        self.conv2d2 = torch.nn.Conv2d(16, 16, 3, 1, 1, bias=True)\n",
    "        self.relu2 = torch.nn.ReLU()\n",
    "        self.conv2d3 = torch.nn.Conv2d(16, 16, 3, 1, 1, bias=True)\n",
    "        self.conv2d4 = torch.nn.Conv2d(16, 16, 3, 1, 1, bias=True)\n",
    "    def forward(self, x):\n",
    "        x = self.conv2d(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.conv2d2(x)\n",
    "        x = self.relu2(x)\n",
    "        x = self.conv2d3(x)\n",
    "        x = self.conv2d4(x)\n",
    "        return x\n",
    "input_shape = [1, 3, 32, 32]\n",
    "torch_model = M().eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "97f3127a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "from tvm import relax\n",
    "from tvm.relax.frontend.torch import from_exported_program\n",
    "from torch.export import export\n",
    "\n",
    "# Give an example argument to torch.export\n",
    "example_args = (torch.randn(1, 3, 32, 32, dtype=torch.float32),)\n",
    "\n",
    "# Convert the model to IRModule\n",
    "with torch.no_grad():\n",
    "    exported_program = export(torch_model, example_args)\n",
    "    run_mod = from_exported_program(exported_program, keep_params_as_input=False)\n",
    "# run_mod, params = relax.frontend.detach_params(run_mod)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f5ca927e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tuple(R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)):\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>dataflow():\n",
       "            lv: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>nn<span style=\"color: #A2F; font-weight: bold\">.</span>conv2d(x, metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">0</span>], strides<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], padding<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], dilation<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], groups<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, data_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NCHW&quot;</span>, kernel_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;OIHW&quot;</span>, out_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NCHW&quot;</span>, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)\n",
       "            lv1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>nn<span style=\"color: #A2F; font-weight: bold\">.</span>relu(lv)\n",
       "            lv2: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>nn<span style=\"color: #A2F; font-weight: bold\">.</span>conv2d(lv1, metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">1</span>], strides<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], padding<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], dilation<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], groups<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, data_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NCHW&quot;</span>, kernel_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;OIHW&quot;</span>, out_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NCHW&quot;</span>, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)\n",
       "            lv3: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>reshape(metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">2</span>], R<span style=\"color: #A2F; font-weight: bold\">.</span>shape([<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>]))\n",
       "            lv4: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>add(lv2, lv3)\n",
       "            lv5: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>nn<span style=\"color: #A2F; font-weight: bold\">.</span>relu(lv4)\n",
       "            lv6: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>nn<span style=\"color: #A2F; font-weight: bold\">.</span>conv2d(lv5, metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">3</span>], strides<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], padding<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], dilation<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], groups<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, data_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NCHW&quot;</span>, kernel_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;OIHW&quot;</span>, out_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NCHW&quot;</span>, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)\n",
       "            lv7: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>reshape(metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">4</span>], R<span style=\"color: #A2F; font-weight: bold\">.</span>shape([<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>]))\n",
       "            lv8: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>add(lv6, lv7)\n",
       "            lv9: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>nn<span style=\"color: #A2F; font-weight: bold\">.</span>conv2d(lv8, metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">5</span>], strides<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], padding<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], dilation<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], groups<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, data_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NCHW&quot;</span>, kernel_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;OIHW&quot;</span>, out_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NCHW&quot;</span>, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)\n",
       "            lv10: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>reshape(metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">6</span>], R<span style=\"color: #A2F; font-weight: bold\">.</span>shape([<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>]))\n",
       "            lv11: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>add(lv9, lv10)\n",
       "            gv: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tuple(R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">=</span> (lv11,)\n",
       "            R<span style=\"color: #A2F; font-weight: bold\">.</span>output(gv)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "\n",
       "<span style=\"color: #007979; font-style: italic\"># Metadata omitted. Use show_meta=True in script() method to show it.</span>\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "run_mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bfd5260e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'tvm.ir.op.Op'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.Var'>\n",
      "<class 'tvm.relax.expr.Constant'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.Call'>\n",
      "<class 'tvm.ir.op.Op'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.DataflowVar'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.Call'>\n",
      "<class 'tvm.ir.op.Op'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.DataflowVar'>\n",
      "<class 'tvm.relax.expr.Constant'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.Call'>\n",
      "<class 'tvm.ir.op.Op'>\n",
      "<class 'tvm.relax.expr.Constant'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.Call'>\n",
      "<class 'tvm.ir.op.Op'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.DataflowVar'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.DataflowVar'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.Call'>\n",
      "<class 'tvm.ir.op.Op'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.DataflowVar'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.Call'>\n",
      "<class 'tvm.ir.op.Op'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.DataflowVar'>\n",
      "<class 'tvm.relax.expr.Constant'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.Call'>\n",
      "<class 'tvm.ir.op.Op'>\n",
      "<class 'tvm.relax.expr.Constant'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.Call'>\n",
      "<class 'tvm.ir.op.Op'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.DataflowVar'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.DataflowVar'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.Call'>\n",
      "<class 'tvm.ir.op.Op'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.DataflowVar'>\n",
      "<class 'tvm.relax.expr.Constant'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.Call'>\n",
      "<class 'tvm.ir.op.Op'>\n",
      "<class 'tvm.relax.expr.Constant'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.Call'>\n",
      "<class 'tvm.ir.op.Op'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.DataflowVar'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.DataflowVar'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.Call'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.DataflowVar'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.Tuple'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.Var'>\n",
      "<class 'tvm.relax.expr.ShapeExpr'>\n",
      "<class 'tvm.relax.expr.SeqExpr'>\n",
      "<class 'tvm.relax.expr.Function'>\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    main <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #008000; font-weight: bold\">None</span>\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import hashlib\n",
    "from tvm.relax.analysis import post_order_visit\n",
    "from tvm.relax import expr as _expr\n",
    "from tvm.ir.op import Op\n",
    "\n",
    "specific_op_names = [\"nn.conv2d\"]\n",
    "specific_op_names = [f\"relax.{name}\" for name in specific_op_names]\n",
    "specific_ops = {}\n",
    "calls = []\n",
    "def fvisit(expr):\n",
    "    print(type(expr))\n",
    "    # if isinstance(expr, Op):\n",
    "    #     if expr.name in specific_op_names:\n",
    "    #         hash_value = hashlib.sha256(tvm.ir.save_json(expr).encode(\"utf-8\")).hexdigest()\n",
    "    #         specific_ops[hash_value] = expr\n",
    "    if isinstance(expr, _expr.Call):\n",
    "        if expr.op.name in specific_op_names:\n",
    "            # expr = expr.args[0]\n",
    "            # hash_value = hashlib.sha256(tvm.ir.save_json(expr).encode(\"utf-8\")).hexdigest()\n",
    "            # specific_ops[hash_value] = expr\n",
    "            calls.append(expr)\n",
    "    # elif isinstance(expr, _expr.Var):\n",
    "    #     hash_value = hashlib.sha256(tvm.ir.save_json(expr).encode(\"utf-8\")).hexdigest()\n",
    "    #     specific_ops[hash_value] = expr\n",
    "    # elif isinstance(expr, _expr.Function):\n",
    "    #     expr = expr.body\n",
    "    #     hash_value = hashlib.sha256(tvm.ir.save_json(expr).encode(\"utf-8\")).hexdigest()\n",
    "    #     specific_ops[hash_value] = expr\n",
    "expr = run_mod[\"main\"]\n",
    "post_order_visit(expr, fvisit)\n",
    "expr = _expr.Tuple([v for v in specific_ops.values()])\n",
    "mod = tvm.IRModule.from_expr(expr)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "de4afd48",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# from tvm.script import relax as R\n",
      "\n",
      "@R.function\n",
      "def main(x: R.Tensor((1, 3, 32, 32), dtype=\"float32\")) -> R.Tuple(R.Tensor((1, 16, 32, 32), dtype=\"float32\")):\n",
      "    with R.dataflow():\n",
      "        lv: R.Tensor((1, 16, 32, 32), dtype=\"float32\") = R.nn.conv2d(x, metadata[\"relax.expr.Constant\"][0], strides=[1, 1], padding=[1, 1, 1, 1], dilation=[1, 1], groups=1, data_layout=\"NCHW\", kernel_layout=\"OIHW\", out_layout=\"NCHW\", out_dtype=\"float32\")\n",
      "        lv1: R.Tensor((1, 16, 32, 32), dtype=\"float32\") = R.nn.relu(lv)\n",
      "        lv2: R.Tensor((1, 16, 32, 32), dtype=\"float32\") = R.nn.conv2d(lv1, metadata[\"relax.expr.Constant\"][1], strides=[1, 1], padding=[1, 1, 1, 1], dilation=[1, 1], groups=1, data_layout=\"NCHW\", kernel_layout=\"OIHW\", out_layout=\"NCHW\", out_dtype=\"float32\")\n",
      "        lv3: R.Tensor((1, 16, 1, 1), dtype=\"float32\") = R.reshape(metadata[\"relax.expr.Constant\"][2], R.shape([1, 16, 1, 1]))\n",
      "        lv4: R.Tensor((1, 16, 32, 32), dtype=\"float32\") = R.add(lv2, lv3)\n",
      "        lv5: R.Tensor((1, 16, 32, 32), dtype=\"float32\") = R.nn.relu(lv4)\n",
      "        lv6: R.Tensor((1, 16, 32, 32), dtype=\"float32\") = R.nn.conv2d(lv5, metadata[\"relax.expr.Constant\"][3], strides=[1, 1], padding=[1, 1, 1, 1], dilation=[1, 1], groups=1, data_layout=\"NCHW\", kernel_layout=\"OIHW\", out_layout=\"NCHW\", out_dtype=\"float32\")\n",
      "        lv7: R.Tensor((1, 16, 1, 1), dtype=\"float32\") = R.reshape(metadata[\"relax.expr.Constant\"][4], R.shape([1, 16, 1, 1]))\n",
      "        lv8: R.Tensor((1, 16, 32, 32), dtype=\"float32\") = R.add(lv6, lv7)\n",
      "        lv9: R.Tensor((1, 16, 32, 32), dtype=\"float32\") = R.nn.conv2d(lv8, metadata[\"relax.expr.Constant\"][5], strides=[1, 1], padding=[1, 1, 1, 1], dilation=[1, 1], groups=1, data_layout=\"NCHW\", kernel_layout=\"OIHW\", out_layout=\"NCHW\", out_dtype=\"float32\")\n",
      "        lv10: R.Tensor((1, 16, 1, 1), dtype=\"float32\") = R.reshape(metadata[\"relax.expr.Constant\"][6], R.shape([1, 16, 1, 1]))\n",
      "        lv11: R.Tensor((1, 16, 32, 32), dtype=\"float32\") = R.add(lv9, lv10)\n",
      "        gv: R.Tuple(R.Tensor((1, 16, 32, 32), dtype=\"float32\")) = (lv11,)\n",
      "        R.output(gv)\n",
      "    return gv\n",
      "\n",
      "# Metadata omitted. Use show_meta=True in script() method to show it.\n"
     ]
    }
   ],
   "source": [
    "print(run_mod[\"main\"])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
